import numpy as np
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt
from python_ai.common.xcommon import sep


def x_pca_impl_again(x, k):
    m = len(x)
    x -= x.mean(axis=0)
    cov = np.dot(x.T, x) / m  # ATTENTION "/ m" for Expection of (x-mu)**2 or (x1-mu1)(x2-mu2)
    lmds, us = np.linalg.eig(cov)
    idx = np.argsort(lmds)
    idx = idx[-1::-1]
    lmds = lmds[idx]
    us = us[idx]
    lmds = lmds[:k]
    us = us[:k]
    x_new = np.dot(x, us.T)
    return x_new, lmds, us


def x_pca_impl_svd_again(x, k):
    m = len(x)
    x -= x.mean(axis=0)
    cov = np.cov(x.T)
    us, lmds, _ = np.linalg.svd(cov)
    us = us[:, :k]  # ATTENTION us[:, :k]
    lmds = lmds[:k]
    x_new = np.dot(x, us)
    return x_new, lmds, us.T  # ATTENTION us.T


def x_cross_check_us(us):
    n_us = len(us)
    for i in range(n_us):
        for j in range(n_us):
            print(np.dot(us[i], us[j]), end=', ')
        print()


if '__main__' == __name__:
    spr = 2
    spc = 4
    spn = 0
    plt.figure(figsize=[16, 8])
    from sklearn.datasets import load_breast_cancer
    x, y = load_breast_cancer(return_X_y=True)
    from sklearn.preprocessing import StandardScaler
    x = StandardScaler().fit_transform(x)

    title = 'pca by PCA => 3 dim'
    sep(title)
    spn += 1
    ax = plt.subplot(spr, spc, spn, projection='3d')
    ax.set_title(title)
    from sklearn.decomposition import PCA
    dc = PCA(n_components=3)
    x_new = dc.fit_transform(x)
    print(dc.explained_variance_)
    x_cross_check_us(dc.components_)
    ax.scatter3D(x_new[:, 0], x_new[:, 1], x_new[:, 2], s=1, c=y)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    title = 'pca by PCA => 2 dim'
    sep(title)
    spn += 1
    ax = plt.subplot(spr, spc, spn)
    ax.set_title(title)
    dc = PCA(n_components=2)
    x_new = dc.fit_transform(x)
    print(dc.explained_variance_)
    x_cross_check_us(dc.components_)
    ax.scatter(x_new[:, 0], x_new[:, 1], s=1, c=y)

    title = 'pca by TruncatedSVD => 3 dim'
    sep(title)
    spn += 1
    ax = plt.subplot(spr, spc, spn, projection='3d')
    ax.set_title(title)
    from sklearn.decomposition import TruncatedSVD
    dc = TruncatedSVD(n_components=3)
    x_new = dc.fit_transform(x)
    print(dc.explained_variance_)
    x_cross_check_us(dc.components_)
    ax.scatter3D(x_new[:, 0], x_new[:, 1], x_new[:, 2], s=1, c=y)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    title = 'pca by TruncatedSVD => 2 dim'
    sep(title)
    spn += 1
    ax = plt.subplot(spr, spc, spn)
    ax.set_title(title)
    dc = TruncatedSVD(n_components=2)
    x_new = dc.fit_transform(x)
    print(dc.explained_variance_)
    x_cross_check_us(dc.components_)
    ax.scatter(x_new[:, 0], x_new[:, 1], s=1, c=y)

    title = 'pca by self impl => 3 dim'
    sep(title)
    spn += 1
    ax = plt.subplot(spr, spc, spn, projection='3d')
    ax.set_title(title)
    x_new, lmds, us = x_pca_impl_again(x, 3)
    print(lmds)
    x_cross_check_us(us)
    ax.scatter3D(x_new[:, 0], x_new[:, 1], x_new[:, 2], s=1, c=y)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    title = 'pca by self impl => 2 dim'
    sep(title)
    spn += 1
    ax = plt.subplot(spr, spc, spn)
    ax.set_title(title)
    x_new, lmds, us = x_pca_impl_again(x, 2)
    print(lmds)
    x_cross_check_us(us)
    ax.scatter(x_new[:, 0], x_new[:, 1], s=1, c=y)

    title = 'pca by self svd impl => 3 dim'
    sep(title)
    spn += 1
    ax = plt.subplot(spr, spc, spn, projection='3d')
    ax.set_title(title)
    x_new, lmds, us = x_pca_impl_svd_again(x, 3)
    print(lmds)
    x_cross_check_us(us)
    ax.scatter3D(x_new[:, 0], x_new[:, 1], x_new[:, 2], s=1, c=y)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    title = 'pca by self svd impl => 2 dim'
    sep(title)
    spn += 1
    ax = plt.subplot(spr, spc, spn)
    ax.set_title(title)
    x_new, lmds, us = x_pca_impl_svd_again(x, 2)
    print(lmds)
    x_cross_check_us(us)
    ax.scatter(x_new[:, 0], x_new[:, 1], s=1, c=y)

    plt.show()
